from typing import Dict, List, Any, Optional
from pydantic import BaseModel

class WorkflowState(BaseModel):
    # 输入
    user_input: str
    
    # 实体识别结果
    entities: Optional[Dict[str, Any]] = None
    
    # 图检索结果
    schema_info: Optional[Dict[str, Any]] = None
    
    # 意图识别结果
    intent: Optional[str] = None
    intent_details: Optional[Dict[str, Any]] = None
    
    # SQL相关
    generated_sql: Optional[str] = None
    validated_sql: Optional[str] = None
    
    # 执行结果
    query_result: Optional[Any] = None
    
    # 最终响应
    final_response: Optional[str] = None
    
    # 错误信息
    error: Optional[str] = None
    
    # 用于兼容原有代码的消息字段
    messages: Optional[List[Any]] = None